#pragma once
#include "../CGIPrefix.h"
#include "CGIVulkanResource.h"
namespace Alice
{
	class CGIVulkanCommandBuffer;
	enum class CGIAccess : uint32;
    struct CGIVulkanPipelineBarrier
    {
    	std::vector<VkImageMemoryBarrier2> ImageBarriers;
		void AddImageLayoutTransition(VkImage Image, VkImageLayout SrcLayout, VkImageLayout DstLayout, const VkImageSubresourceRange& SubresourceRange);
    	void Execute(CGIVulkanCommandBuffer* CmdBuffer);
    	void Execute(VkCommandBuffer CmdBuffer);
        static VkImageSubresourceRange MakeSubresourceRange(VkImageAspectFlags AspectMask, uint32 FirstMip = 0, uint32 NumMips = VK_REMAINING_MIP_LEVELS, uint32 FirstLayer = 0, uint32 NumLayers = VK_REMAINING_ARRAY_LAYERS);
    };
    struct VulkanImageLayout
	{
		VulkanImageLayout(VkImageLayout InitialLayout, uint32 InNumMips, uint32 InNumLayers, VkImageAspectFlags Aspect) :
			NumMips(InNumMips),
			NumLayers(InNumLayers),
			NumPlanes((Aspect == (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) ? 2 : 1),
			MainLayout(InitialLayout)
		{
		}
    	VulkanImageLayout(const VulkanImageLayout& inOther);
		uint32 NumMips;
		uint32 NumLayers;
		uint32 NumPlanes;
		// The layout when all the subresources are in the same state.
		VkImageLayout MainLayout;
		// Explicit subresource layouts. Always NumLayers*NumMips elements.
		std::vector<VkImageLayout> SubresLayouts;

		inline bool AreAllSubresourcesSameLayout() const
		{
			return SubresLayouts.size() == 0;
		}

		VkImageLayout GetSubresLayout(uint32 Layer, uint32 Mip, VkImageAspectFlagBits Aspect) const
		{
			return GetSubresLayout(Layer, Mip, (Aspect==VK_IMAGE_ASPECT_STENCIL_BIT) ? NumPlanes - 1 : 0);
		}

		VkImageLayout GetSubresLayout(uint32 Layer, uint32 Mip, uint32 Plane) const
		{
			if (SubresLayouts.size() == 0)
			{
				return MainLayout;
			}

			if (Layer == (uint32)-1)
			{
				Layer = 0;
			}
			return SubresLayouts[(Plane * NumLayers * NumMips) + (Layer * NumMips) + Mip];
		}

		bool AreSubresourcesSameLayout(VkImageLayout Layout, const VkImageSubresourceRange& SubresourceRange) const;

		inline uint32 GetSubresRangeLayerCount(const VkImageSubresourceRange& SubresourceRange) const
		{
			return (SubresourceRange.layerCount == VK_REMAINING_ARRAY_LAYERS) ? (NumLayers - SubresourceRange.baseArrayLayer) : SubresourceRange.layerCount;
		}

		inline uint32 GetSubresRangeMipCount(const VkImageSubresourceRange& SubresourceRange) const
		{
			return (SubresourceRange.levelCount == VK_REMAINING_MIP_LEVELS) ? (NumMips - SubresourceRange.baseMipLevel) : SubresourceRange.levelCount;
		}

		void CollapseSubresLayoutsIfSame();

		void Set(VkImageLayout Layout, const VkImageSubresourceRange& SubresourceRange);
	};
	void VulkanSetImageLayout(CGIVulkanCommandBuffer* CmdBuffer, VkImage Image, VkImageLayout OldLayout, VkImageLayout NewLayout, const VkImageSubresourceRange& SubresourceRange);
	class VulkanLayoutManager
	{
	public:
		VulkanLayoutManager(bool InWriteOnly, VulkanLayoutManager* InFallback)
			: bWriteOnly(InWriteOnly)
			, Fallback(InFallback)
		{
		}
		
		void NotifyDeletedImage(VkImage Image);

		// Predetermined layouts for given RHIAccess
		static VkImageLayout GetDefaultLayout(CGIVulkanCommandBuffer* CmdBuffer, const CGIVulkanTexture& VulkanTexture, CGIAccess DesiredAccess);

		// Expected layouts and Hints are workarounds until we can use 'hardcoded layouts' everywhere.
		static VkImageLayout SetExpectedLayout(CGIVulkanCommandBuffer* CmdBuffer, const CGIVulkanTexture& VulkanTexture, CGIAccess DesiredAccess);
		VkImageLayout GetDepthStencilHint(const CGIVulkanTexture& VulkanTexture, VkImageAspectFlagBits AspectBit);

		const VulkanImageLayout* GetFullLayout(VkImage Image) const
		{
			auto iter=Layouts.find(Image);
			const VulkanImageLayout* Layout = iter->second;
			if (!Layout && Fallback)
			{
				return Fallback->GetFullLayout(Image);
			}
			return Layout;
		}

		const VulkanImageLayout* GetFullLayout(const CGIVulkanTexture& VulkanTexture, bool bAddIfNotFound = false, VkImageLayout LayoutIfNotFound = VK_IMAGE_LAYOUT_UNDEFINED)
		{
			auto iter=Layouts.find(VulkanTexture.mImage);
			const VulkanImageLayout* Layout = iter->second;
			
			if (!Layout && Fallback)
			{
				Layout = Fallback->GetFullLayout(VulkanTexture, false);

				// If the layout was found in the fallback, carry it forward to our current manager for future tracking
				if (Layout)
				{
					//Layouts.insert(std::pair<VkImage,VulkanImageLayout*>(VulkanTexture.mImage,Layout));
					return Layout;
				}
			}
			
			if (Layout)
			{
				return Layout;
			}
			else if (!bAddIfNotFound)
			{
				return nullptr;
			}
			VulkanImageLayout *newLayout=new VulkanImageLayout(LayoutIfNotFound, VulkanTexture.GetNumMips(), VulkanTexture.GetNumberOfArrayLevels(), VulkanTexture.GetFullAspectMask());
			Layouts.insert(std::pair<VkImage,VulkanImageLayout*>(VulkanTexture.mImage,newLayout));
			return Layouts[VulkanTexture.mImage];
		}

		// Not the preferred path because we can't ensure Mip and Layer counts match, but still necessary for images like the backbuffer
		void SetFullLayout(VkImage Image, VulkanImageLayout& NewLayout)
		{
			auto iter=Layouts.find(Image);
			VulkanImageLayout* Layout = iter->second;
			if (Layout)
			{
				*Layout = NewLayout;
			}
			else
			{
				Layouts.insert(std::pair<VkImage,VulkanImageLayout*>(Image, &NewLayout));
			}
		}

		void SetFullLayout(const CGIVulkanTexture& VulkanTexture, VulkanImageLayout& NewLayout)
		{
			SetFullLayout(VulkanTexture.mImage, NewLayout);
		}

		void SetFullLayout(const CGIVulkanTexture& VulkanTexture, VkImageLayout InLayout, bool bOnlyIfNotFound=false)
		{
			auto iter=Layouts.find(VulkanTexture.mImage);
			VulkanImageLayout* Layout = NULL;
			if (iter != Layouts.end()) {
				Layout=iter->second;
			}
			if (Layout)
			{
				if (!bOnlyIfNotFound)
				{
					Layout->Set(InLayout, CGIVulkanPipelineBarrier::MakeSubresourceRange(VulkanTexture.GetFullAspectMask()));
				}
			}
			else
			{
				VulkanImageLayout *newLayout=new VulkanImageLayout(InLayout, VulkanTexture.GetNumMips(), VulkanTexture.GetNumberOfArrayLevels(), VulkanTexture.GetFullAspectMask());
				Layouts.insert(std::pair<VkImage,VulkanImageLayout*>(VulkanTexture.mImage,newLayout));
			}
		}

		static bool IsPartialResource(const CGIVulkanTexture& VulkanTexture, const VkImageSubresourceRange& InSubresourceRange)
		{
			return (VulkanTexture.GetFullAspectMask() != InSubresourceRange.aspectMask) ||
				(InSubresourceRange.baseArrayLayer != 0) ||
				(InSubresourceRange.baseMipLevel != 0) ||
				((InSubresourceRange.levelCount != VK_REMAINING_MIP_LEVELS) && (InSubresourceRange.levelCount != VulkanTexture.GetNumMips())) ||
				((InSubresourceRange.layerCount != VK_REMAINING_ARRAY_LAYERS) && (InSubresourceRange.layerCount != VulkanTexture.GetNumberOfArrayLevels()));
		}

		void SetLayout(const CGIVulkanTexture& VulkanTexture, const VkImageSubresourceRange& InSubresourceRange, VkImageLayout InLayout)
		{
			auto iter=Layouts.find(VulkanTexture.mImage);
			VulkanImageLayout* Layout = iter->second;

			// If we're not going to overwrite the entire resource, start from its last known state on the queue
			if (!Layout && IsPartialResource(VulkanTexture, InSubresourceRange) && Fallback)
			{
				const VulkanImageLayout* FallbackLayout = Fallback->GetFullLayout(VulkanTexture, false);
				if (FallbackLayout)
				{
					//Layouts.insert(std::pair<VkImage,VulkanImageLayout>(VulkanTexture.mImage,*FallbackLayout));
					//Layout = &Layouts[VulkanTexture.mImage];
				}
			}

			if (Layout)
			{
				Layout->Set(InLayout, InSubresourceRange);
			}
			else
			{
				VulkanImageLayout *NewLayout=new VulkanImageLayout(VK_IMAGE_LAYOUT_UNDEFINED, VulkanTexture.GetNumMips(), VulkanTexture.GetNumberOfArrayLevels(), VulkanTexture.GetFullAspectMask());
				NewLayout->Set(InLayout, InSubresourceRange);
				Layouts.insert(std::pair<VkImage,VulkanImageLayout*>(VulkanTexture.mImage, NewLayout));
			}
		}

		// Transfers our layouts into the destination
		void TransferTo(VulkanLayoutManager& Destination);

	private:

		std::unordered_map<VkImage, VulkanImageLayout*> Layouts;

		// If we're WriteOnly, we should never read layout from this instance.  This is important for parallel rendering.
		// When WriteOnly, this instance of the layout manager should only collect layouts to later feed them to the another central mgr.
		const bool bWriteOnly;

		// If parallel command list creation is NOT supported, then the queue's layout mgr can be used as a fallback to fetch previous layouts.
		VulkanLayoutManager* Fallback;
	};
}